package com.flute.haflute.jobbox.base.storage;

import static me.prettyprint.cassandra.utils.StringUtils.string;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import me.prettyprint.cassandra.service.CassandraClient;
import me.prettyprint.cassandra.service.CassandraClientPool;
import me.prettyprint.cassandra.service.CassandraClientPoolFactory;
import me.prettyprint.cassandra.service.Keyspace;

import org.apache.cassandra.thrift.Column;
import org.apache.cassandra.thrift.ColumnParent;
import org.apache.cassandra.thrift.KeyRange;
import org.apache.cassandra.thrift.NotFoundException;
import org.apache.cassandra.thrift.SlicePredicate;
import org.apache.cassandra.thrift.SliceRange;
import org.apache.thrift.TException;

public class RangeSelector {
	private String keySpaceName = "Keyspace1";
	private String cfName = "Standard1";
	private int oneRoundCount = 10;
	private long maxCount = 100;
	private long queriedCount = 0;
	private boolean hasMore = true;
	private CassandraClientPool pool;
	private CassandraClient client;
	private SlicePredicate sp;
	private KeyRange keyRange;
	private ColumnParent columnParent;
	private Keyspace keyspace;
	private String firstKey;
	private String lastKey;
	private String fetchStartKey;
	private String fetchEndKey;

	public RangeSelector(String cassandraServer, String keySpace,
			String columnFamily) throws Exception {
		this(cassandraServer);
		this.keySpaceName = keySpace;
		this.cfName = columnFamily;
	}

	public RangeSelector(String cassandraServer) throws Exception {
		pool = CassandraClientPoolFactory.INSTANCE.get();
		client = pool.borrowClient(cassandraServer);
		initDefaultPredicate();
	}

	public RangeSelector(CassandraClient client) {
		this.client = client;
		initDefaultPredicate();
	}

	private void initDefaultPredicate() {
		sp = new SlicePredicate();
		SliceRange sr = new SliceRange();
		sr.setStart(new byte[0]);
		sr.setFinish(new byte[0]);
		sp.setSlice_range(sr);
	}

	private void initQueryPredicate(String startKey, String endKey, int count)
			throws IllegalArgumentException, NotFoundException, TException {
		keyspace = client.getKeyspace(keySpaceName);
		columnParent = new ColumnParent(cfName);

		firstKey = "";
		char ch = 127;
		lastKey = String.valueOf(ch);
		if (startKey != null)
			firstKey = startKey;
		if (endKey != null)
			lastKey = endKey;
		fetchStartKey = firstKey;
		fetchEndKey = lastKey;
		keyRange = new KeyRange();
		keyRange.setCount(count);
		keyRange.setStart_key(firstKey.getBytes());
		keyRange.setEnd_key(lastKey.getBytes());
	}

	public void setFetchCount(int count) {
		this.oneRoundCount = count;
	}

	public void prepareQueryStatement(String startKey, String endKey,
			long maxCount, int fetchCount) throws Exception {
		if (maxCount > fetchCount) {
			this.maxCount = maxCount;
		} else {
			this.maxCount = Long.MAX_VALUE;
		}

		initQueryPredicate(startKey, endKey, fetchCount);
	}

	public Map<String, Map<String, String>> fetchNext() throws Exception {
		return fetchNext(oneRoundCount);
	}

	public boolean hasNext() {
		return hasMore;
	}

	public Map<String, Map<String, String>> fetchNext(int count)
			throws Exception {
		if (!hasMore) {
			return null;
		}

		int realCount = count;
		queriedCount += realCount;
		boolean needFirst = true;
		if (!fetchStartKey.equals(firstKey)) {
			realCount++;
			needFirst = false;
		}
		keyRange.setCount(realCount);
		// may have one dup
		keyRange.setStart_key(fetchStartKey.getBytes());
		keyRange.setEnd_key(fetchEndKey.getBytes());
		Map<String, List<Column>> rangeResult = keyspace.getRangeSlices(
				columnParent, sp, keyRange);
		if (!needFirst) {
			rangeResult.remove(fetchStartKey);
		}

		Map<String, Map<String, String>> result = new HashMap<String, Map<String, String>>(
				rangeResult.size());
		for (Entry<String, List<Column>> entry : rangeResult.entrySet()) {
			fetchStartKey = entry.getKey();
			Map<String, String> columns = new HashMap<String, String>(entry
					.getValue().size());
			for (Column col : entry.getValue()) {
				columns.put(string(col.getName()), string(col.getValue()));
			}
			result.put(fetchStartKey, columns);
		}

		hasMore = rangeResult.size() == count && !fetchStartKey.equals(lastKey)
				&& queriedCount < maxCount;

		return result;
	}

	public long getTotalCount() throws Exception {
		if (keyspace == null)
			throw new Exception("database not connected");

		SlicePredicate countSP = new SlicePredicate();
		SliceRange sr = new SliceRange();
		sr.setStart(new byte[0]);
		sr.setFinish(new byte[100]);
		countSP.setSlice_range(sr);

		KeyRange countKR = new KeyRange();
		countKR.setCount(maxCount > 1000 ? 1000 : (int) maxCount);
		countKR.setStart_key(firstKey.getBytes());
		countKR.setEnd_key(lastKey.getBytes());

		Map<String, List<Column>> rangeResult = keyspace.getRangeSlices(
				columnParent, countSP, countKR);

		if (rangeResult.size() == 0)
			return 0;

		long count = 1;
		while (rangeResult.size() > 1) {
			count += rangeResult.size() - 1;
			String newStartKey = firstKey;
			Iterator<String> ite = rangeResult.keySet().iterator();
			while (ite.hasNext()) {
				newStartKey = ite.next();
			}
			countKR.setStart_key(newStartKey.getBytes());
			long moreCount = maxCount - count;
			if (moreCount <= 0) {
				break;
			}
			countKR.setCount(moreCount > 1000 ? 1000 : (int) moreCount);
			rangeResult = keyspace.getRangeSlices(columnParent, countSP,
					countKR);
		}

		return count;
	}

	public void close() {
		if (pool != null) {
			try {
				pool.releaseClient(client);
				keyspace = null;
			} catch (Exception e) {
				e.printStackTrace();
			}
		}
	}

	/**
	 * @param args
	 */
	public static void main(String[] args) {
		try {
			System.err.println("Usage:");
			System.err
					.println("connect/C host:port    eg. C 192.168.1.128:9160");
			System.err
					.println("select s=startKey e=endKey t=totalKeys p=keysPerPage       eg. select s=a e=zz t=100 p=10");
			System.err.println("quit to exit");
			boolean console = false;
			String cassandraServer = "192.168.1.128:9160";
			String startKey = "";
			String endKey = "";
			long total = -1;
			int each = 10;
			if (args.length == 5) {
				cassandraServer = args[0];
				startKey = args[1];
				endKey = args[2];
				total = Long.parseLong(args[3]);
				each = Integer.parseInt(args[4]);
			} else if (args.length == 1 && args[0].equalsIgnoreCase("cmd")) {
				console = true;
			}

			if (console) {
				BufferedReader reader = new BufferedReader(
						new InputStreamReader(System.in));
				String cmd;
				RangeSelector selector = null;
				System.err.print("command>");
				while ((cmd = reader.readLine()) != null) {
					if (cmd.startsWith("connect ") || cmd.startsWith("C ")) {
						cassandraServer = cmd.replace("connect ", "").replace(
								"C ", "");
						selector = new RangeSelector(cassandraServer);
					} else if (cmd.equalsIgnoreCase("quit")) {
						if (selector != null) {
							selector.close();
						}
						System.exit(0);
					} else if (cmd.startsWith("select")) {
						cmd = cmd.replace("select", "");
						String[] params = cmd.split(" ");
						for (String s : params) {
							if (s.startsWith("s="))
								startKey = s.replace("s=", "");
							else if (s.startsWith("e="))
								endKey = s.replace("e=", "");
							else if (s.startsWith("t="))
								total = Long.parseLong(s.replace("t=", ""));
							else if (s.startsWith("p="))
								each = Integer.parseInt(s.replace("p=", ""));
						}

						if (selector != null) {
							System.out.println("now querying, please wait...");
							selector.prepareQueryStatement(startKey, endKey,
									total, each);
							System.err.println("=== total:"
									+ selector.getTotalCount());
							Map<String, Map<String, String>> qr = selector
									.fetchNext(each);
							for (Entry<String, Map<String, String>> entry : qr
									.entrySet()) {
								System.out.print("key=" + entry.getKey());
								for (Entry<String, String> colEntry : entry
										.getValue().entrySet()) {
									// System.out.print("\tcolName=" +
									// colEntry.getKey());
									System.out.println("\tValue="
											+ colEntry.getValue());
								}
							}
							System.out
									.println("-----------------------------------------More ------------------------------------\n press M to get more");
						} // if(select != null
					} else if (cmd.equals(" ") || cmd.equalsIgnoreCase("m")) {
						if (selector != null && selector.hasNext()) {
							Map<String, Map<String, String>> qr = selector
									.fetchNext(each);
							for (Entry<String, Map<String, String>> entry : qr
									.entrySet()) {
								System.out.print("key=" + entry.getKey());
								for (Entry<String, String> colEntry : entry
										.getValue().entrySet()) {
									// System.out.print("\tcolName=" +
									// colEntry.getKey());
									System.out.println("\tValue="
											+ colEntry.getValue());
								}
							}
						}
						System.out
								.println("-----------------------------------------More ------------------------------------\n press M to get more");
					}
					System.err.print("command>");
				} // while
			} else {
				RangeSelector selector = new RangeSelector(cassandraServer);
				selector.prepareQueryStatement(startKey, endKey, total, each);
				System.out.println("total:" + selector.getTotalCount());
				while (selector.hasNext()) {
					Map<String, Map<String, String>> qr = selector.fetchNext();
					for (Entry<String, Map<String, String>> entry : qr
							.entrySet()) {
						System.out.print("key=" + entry.getKey());
						for (Entry<String, String> colEntry : entry.getValue()
								.entrySet()) {
							System.out.print("\tcolName=" + colEntry.getKey());
							System.out
									.println("\tValue=" + colEntry.getValue());
						}
					}
				}
				selector.close();
			}

		} catch (Exception e) {
			e.printStackTrace();
		}

	}

}
